﻿using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Reflection;
using System.Diagnostics.Contracts;
using System.Collections;

using HeapModel.ProgInfo;
using HeapModel.ConcreteMemory;

namespace HeapProfiler.Runtime
{
	/// <summary>
	/// Allows to build heaps
	/// </summary>
	public sealed class ConcreteMemoryBuilder
	{
		//keep track of visited objects and the graph constructed
		private ConcreteHeap m_heap;
		private readonly Dictionary<Object, ConcreteNode> m_visited;
		private int m_heapctr;
		private readonly HashSet<Type> m_initializedtypeset;
		private readonly IdentifierIntern m_interntable;

		private int m_accesspathtokenctr;
		private readonly Dictionary<OffsetIdentifier, AccessPath> m_singlestepaccesspath;
		private readonly Dictionary<Tuple<String, OffsetIdentifier>, AccessPath> m_accesspathextend;

		/// <summary>
		/// The colletion of type prefixes to ignore
		/// </summary>
		public static HashSet<String> IgnoredPrefixes;

		static ConcreteMemoryBuilder()
		{
			IgnoredPrefixes = new HashSet<String>()
            {
                "Microsoft", 
                "PresentationFramework", 
                "PresentationCore", 
                "WindowsBase", 
                "Assembly", 
                "vshost32", 
                "FXAssembly", 
                "System", 
                "mscorlib", 
                "HeapModel",
 				"HeapProfiler"
            };
		}

		/// <summary>
		/// Initializes an empty heap
		/// </summary>
		public ConcreteMemoryBuilder()
		{
			//initialize extraction information
			this.m_heap = new ConcreteHeap();
			this.m_visited = new Dictionary<Object, ConcreteNode>(new StrictRefEqualityCompare());
			this.m_heapctr = 1;

			this.m_initializedtypeset = new HashSet<Type>();
			this.m_interntable = new IdentifierIntern();

			this.m_accesspathtokenctr = 1;
			this.m_singlestepaccesspath = new Dictionary<OffsetIdentifier, AccessPath>();
			this.m_accesspathextend = new Dictionary<Tuple<String, OffsetIdentifier>, AccessPath>();
		}

		/// <summary>
		/// Reset the concrete heap information before taking a new snapshot
		/// </summary>
		public void ResetConcreteHeap()
		{
			this.m_heap = new ConcreteHeap();
			this.m_visited.Clear();
			this.m_heapctr = 1;
		}

		/// <summary>
		/// Strict reference equality comparision for visited dictionary.
		/// </summary>
		sealed class StrictRefEqualityCompare : IEqualityComparer<Object>
		{
			bool IEqualityComparer<object>.Equals(object x, object y)
			{ return Object.ReferenceEquals(x, y); }

			int IEqualityComparer<object>.GetHashCode(object obj)
			{ return (obj == null ? 0 : obj.GetHashCode()); }
		}

		/// <summary>
		/// Walk and extact a concrete heap model.
		/// </summary>
		public ConcreteHeap ExtractConcreteHeap(int paramc, Object[] paramarray, bool[] samplemask, String[] names)
		{
			this.AddStackRoots(paramc, paramarray, samplemask, names);
			this.AddStaticRoots();

			this.m_interntable.CompleteOffsetInfo();

			return this.m_heap;
		}

		/// <summary>
		/// Gets the current heap snapshot.
		/// </summary>
		public ConcreteHeap Heap
		{ get { return this.m_heap; } }

		/// <summary>
		/// Gets all the type names that appear in the current heap snapshot for output.
		/// </summary>
		public HashSet<TypeIdentifier> AllTypesSet
		{ get { return new HashSet<TypeIdentifier>(this.m_interntable.AllTypesForDump()); } }

		/// <summary>
		/// Gets the offsets that appear in the current heap snapshot for output.
		/// </summary>
		public HashSet<OffsetIdentifier> AllOffsetsSet
		{ get { return new HashSet<OffsetIdentifier>(this.m_interntable.AllFieldForDump()); } }

		/// <summary>
		/// Add roots collected from the stack
		/// </summary>
		public void AddStackRoots(int paramc, Object[] paramarray, bool[] samplemask, String[] names)
		{
			Contract.Requires(paramc > 0);
			Contract.Requires(paramarray != null);
			Contract.Requires(samplemask != null);
			Contract.Requires(names != null);

			//get the roots from the stack
			for(int i = 0; i < paramc; ++i)
			{
				bool dosample = samplemask[i];
				String parameterName = names[i];
				Object parameter = paramarray[i];

				this.TryAddRoot(dosample, parameterName, parameter, true);
			}
		}

		/// <summary>
		/// Tries to add a single root to the heap graph
		/// </summary>
		public bool TryAddRoot(bool dosample, String name, Object root, bool isstackroot)
		{
			if(dosample)
			{
				ConcreteNode trgt = VisitObject(root, null, null);
				this.m_heap.AddSingleRootTarget(name, isstackroot, trgt);
				return true;
			}

			return false;
		}

		/// <summary>
		/// Walks the static field roots.
		/// </summary>
		public void AddStaticRoots()
		{
			//get all the static roots
			var asmarr = AppDomain.CurrentDomain.GetAssemblies();
			var badnames = IgnoredPrefixes;
			foreach(var asm in asmarr)
			{
				var fns = asm.FullName;
				if(badnames.Any((String s) => fns.StartsWith(s)))
					continue;

				var typearr = asm.GetTypes();
				foreach(Type tt in typearr)
				{
					if(tt.Namespace != null && badnames.Any((String s) => tt.Namespace.StartsWith(s)))
						continue;

					var sff = tt.GetFields(BindingFlags.Public | BindingFlags.NonPublic | BindingFlags.Static);
					foreach(var sf in sff)
					{
						if(IdentifierIntern.IsStructType(sf.FieldType) || IdentifierIntern.IsObjectType(sf.FieldType))
						{
							TypeIdentifier fty = this.m_interntable.InternTypeName(sf.FieldType);
							String sfn = IdentifierIntern.NiceTypeNameString(sf.DeclaringType) + "." + sf.Name;
							ConcreteNode cobj = VisitObject(sf.GetValue(null), null, null);
							this.m_heap.AddSingleRootTarget(sfn, false, cobj);
						}
					}
				}
			}
		}

		/// <summary>
		/// Intern the access path consisting of the single step given by the offset identifier.
		/// </summary>
		private AccessPath InternSingleOffsetPath(OffsetIdentifier offid)
		{
			Contract.Requires(offid != null);
			Contract.Ensures(Contract.Result<AccessPath>() != null);

			if(this.m_singlestepaccesspath.ContainsKey(offid))
				return this.m_singlestepaccesspath[offid];
			else
			{
				AccessPath naccp = new AccessPath(new OffsetIdentifier[] { offid }, this.m_accesspathtokenctr++);
				this.m_singlestepaccesspath.Add(offid, naccp);
				return naccp;
			}
		}

		/// <summary>
		/// Intern the access path consisting of the single step given by the offset identifier.
		/// </summary>
		private AccessPath ExtendOffsetPath(AccessPath current, OffsetIdentifier offid)
		{
			Contract.Requires(offid != null);
			Contract.Ensures(Contract.Result<AccessPath>() != null);

			Tuple<String, OffsetIdentifier> extkey = new Tuple<String, OffsetIdentifier>(current.PathString, offid);
			if(this.m_accesspathextend.ContainsKey(extkey))
				return this.m_accesspathextend[extkey];
			else
			{
				OffsetIdentifier[] oap = new OffsetIdentifier[current.PathValue.Length + 1];
				Array.Copy(current.PathValue, oap, current.PathValue.Length);
				oap[oap.Length - 1] = offid;
				AccessPath naccp = new AccessPath(oap, this.m_accesspathtokenctr++);

				this.m_accesspathextend.Add(extkey, naccp);
				return naccp;
			}
		}

		/// <summary>
		/// Visit the object (an all reachable memory as needed) returning the node that corresponds to it.
		/// This handles the memoization stuff <c>ProcessObject</c> does the traversal.
		/// </summary>
		private ConcreteNode VisitObject(Object o, AccessPath nestpath, List<ConcretePointer> insidecpl)
		{
			if(o == null)
				return this.m_heap.NullTarget;
			else if(insidecpl == null && this.m_visited.ContainsKey(o))
				return this.m_visited[o];
			else
			{
				Type rawType = o.GetType();
				if(!this.m_initializedtypeset.Contains(rawType))
				{
					this.m_interntable.InternTypeName(rawType);
					this.m_initializedtypeset.Add(rawType);
				}

				return this.ProcessObject(o, nestpath, insidecpl);
			}
		}

		/// <summary>
		/// Process the object (recursively as needed) returning the node that corresponds to it.
		/// This handles the traversal stuff <c>VisitObject</c> does the traversal.
		/// </summary>
		private ConcreteNode ProcessObject(Object o, AccessPath nestpath, List<ConcretePointer> insidecpl)
		{
			Contract.Requires(o != null);
			
			TypeIdentifier tai = this.m_interntable.LookupInternedTypeName(o.GetType());

			if(o is String)
				return this.ProcessStringEntry(o);
			else if(tai.NameValue.StartsWith("System.Reflection"))
				return this.ProcessReflectionEntry(o);
			else if(tai.IsCollectionType)
				return this.ProcessCollectionObject(o);
			else
				return this.ProcessTransparentObjectOrStruct(o, nestpath, insidecpl);
		}

		/// <summary>
		/// Process an object that is representing a collection.
		/// </summary>
		private ConcreteNode ProcessCollectionObject(Object o)
		{
			Contract.Requires(o != null);
			
			Type collty = o.GetType();
			TypeIdentifier tai = this.m_interntable.LookupInternedTypeName(collty);
			int count;
			int datasize;
			int overhead;
			Type contentstype;
			bool processcontents = false;
			OffsetIdentifier offset = null;

			if(tai.IsArrayType)
			{
				Array tarray = (Array)o;
				count = tarray.Length;
				contentstype = tarray.GetType().GetElementType();

				if(contentstype.IsPrimitive | IdentifierIntern.IsStructType(contentstype))
					datasize = count * this.m_interntable.TypeDataSizeMap[new TypeInfoIdentPair(contentstype)];
				else
					datasize = count * 4;
			
				overhead = TypeIdentifier.OBJECT_HEADER_SIZE + TypeIdentifier.ARRAY_SIZE_OVERHEAD;
			}
			else if(tai.IsListType | tai.IsSetType)
			{
				MethodInfo[] miia = o.GetType().GetMethods();
				MethodInfo mienum = miia.FirstOrDefault((MethodInfo mi) => mi.Name.Equals("GetEnumerator"));
				PropertyInfo pcurrent = mienum.ReturnType.GetProperty("Current");

				contentstype = pcurrent.PropertyType;
				if(tai.IsListType)
				{
					count = ((IList)o).Count;
					int underlysize = Math.Max(TypeIdentifier.LIST_MIN_CAPCITY, count);

					if(contentstype.IsPrimitive | IdentifierIntern.IsStructType(contentstype))
						datasize = underlysize * this.m_interntable.TypeDataSizeMap[new TypeInfoIdentPair(contentstype)];
					else
						datasize = underlysize * 4;

					overhead = TypeIdentifier.OBJECT_HEADER_SIZE + TypeIdentifier.LIST_SIZE_OVERHEAD;
				}
				else
				{
					count = 0;
					foreach(Object eo in ((IEnumerable)o))
						++count;

					int underlysize = Math.Max(TypeIdentifier.HASHSET_MIN_CAPCITY, count);

					if(contentstype.IsPrimitive | IdentifierIntern.IsStructType(contentstype))
						datasize = underlysize * TypeIdentifier.HASHSET_SIZE_CAPACITY_FACTOR * this.m_interntable.TypeDataSizeMap[new TypeInfoIdentPair(contentstype)];
					else
						datasize = underlysize * TypeIdentifier.HASHSET_SIZE_CAPACITY_FACTOR * 4;
					
					overhead = TypeIdentifier.OBJECT_HEADER_SIZE + TypeIdentifier.HASHSET_SIZE_OVERHEAD;
				}
			}
			else if(tai.IsDictionaryType)
			{
				IDictionary tdict = (IDictionary)o;
				MethodInfo mienum = o.GetType().GetMethod("GetEnumerator");
				PropertyInfo pcurrent = mienum.ReturnType.GetProperty("Current");

				contentstype = pcurrent.PropertyType;
				count = tdict.Count;
				int underlysize = Math.Max(TypeIdentifier.DICTIONARY_MIN_CAPCITY, count);
				datasize = underlysize * TypeIdentifier.DICTIONARY_SIZE_CAPACITY_FACTOR * this.m_interntable.TypeDataSizeMap[new TypeInfoIdentPair(contentstype)];
				overhead = TypeIdentifier.OBJECT_HEADER_SIZE + TypeIdentifier.DICTIONARY_SIZE_OVERHEAD;
			}
			else
				throw new InvalidOperationException("I Don't know this container type??");

			if(IdentifierIntern.IsObjectType(contentstype) || IdentifierIntern.IsStructType(contentstype))
			{
				processcontents = true;
				offset = this.m_interntable.LookupInternedCollectionFieldName(o.GetType());
			}

			ConcreteNode rr = new ConcreteNode(this.m_heapctr++, tai, count, datasize, overhead);
			this.m_heap.AddNewObject(rr);
			this.m_visited.Add(o, rr);

			List<ConcretePointer> tl = new List<ConcretePointer>();
			if(processcontents)
			{
				AccessPath accp = this.InternSingleOffsetPath(offset);
				bool isstructcontents = IdentifierIntern.IsStructType(contentstype);

				foreach(Object ao in ((System.Collections.IEnumerable)o))
				{
					if(isstructcontents)
						VisitObject(ao, accp, tl);
					else
					{
						ConcreteNode tobj = VisitObject(ao, null, null);
						ConcretePointer cfte = new ConcretePointer(accp, tobj);
						tl.Add(cfte);
					}
				}
			}

			rr.FinalizeTargetSet(tl);
			return rr;
		}

		/// <summary>
		/// Process an object that is representing an class that we want to show all the explicit debug information for.
		/// </summary>
		private ConcreteNode ProcessTransparentObjectOrStruct(Object o, AccessPath nestpath, List<ConcretePointer> insidecpl)
		{
			Contract.Requires(o != null);
			
			Type tty = o.GetType();
			TypeIdentifier tai = this.m_interntable.LookupInternedTypeName(tty);
			int sizeinbytes = this.m_interntable.TypeDataSizeMap[new TypeInfoIdentPair(tty)];
			int overheadsize = IdentifierIntern.IsObjectType(tty) ? TypeIdentifier.OBJECT_HEADER_SIZE : 0;

			ConcreteNode rr;
			List<ConcretePointer> cpl;
			if(insidecpl != null)
			{
				rr = null;
				cpl = insidecpl;
			}
			else
			{
				rr = new ConcreteNode(this.m_heapctr++, tai, sizeinbytes, overheadsize);
				this.m_heap.AddNewObject(rr);
				this.m_visited.Add(o, rr);

				cpl = new List<ConcretePointer>();
			}

			List<FieldInfo> fom = this.m_interntable.InternTypeToRefFieldMap[new TypeInfoIdentPair(tty)];
			foreach(FieldInfo fi in fom)
			{
				if(IdentifierIntern.IsObjectType(fi.FieldType) || IdentifierIntern.IsStructType(fi.FieldType))
				{
					Object tobj = fi.GetValue(o);
					OffsetIdentifier fname = this.m_interntable.LookupInternedFieldName(fi);

					AccessPath extaccp;
					if(nestpath == null)
						extaccp = this.InternSingleOffsetPath(fname);
					else
						extaccp = this.ExtendOffsetPath(nestpath, fname);

					if(IdentifierIntern.IsStructType(fi.FieldType))
						this.VisitObject(tobj, extaccp, cpl);
					else
					{
						ConcreteNode ctrgt = VisitObject(tobj, null, null);
						cpl.Add(new ConcretePointer(extaccp, ctrgt));
					}
				}
			}

			if(rr != null)
				rr.FinalizeTargetSet(cpl);

			return rr;
		}

		/// <summary>
		/// Process an object that is a string.
		/// </summary>
		private ConcreteNode ProcessStringEntry(Object o)
		{
			Contract.Requires(o != null);

			Type tty = o.GetType();
			TypeIdentifier tai = this.m_interntable.LookupInternedTypeName(tty);
			ConcreteNode rr = new ConcreteNode(this.m_heapctr++, tai, ((String)o).Length, TypeIdentifier.OBJECT_HEADER_SIZE + 8);
			this.m_heap.AddNewObject(rr);
			this.m_visited.Add(o, rr);
			rr.FinalizeTargetSet(new List<ConcretePointer>());

			return rr;
		}

		/// <summary>
		/// Process an object that is something from the system reflection namespace.
		/// </summary>
		private ConcreteNode ProcessReflectionEntry(Object o)
		{
			Contract.Requires(o != null);

			Type tty = o.GetType();
			TypeIdentifier tai = this.m_interntable.LookupInternedTypeName(tty);
			ConcreteNode rr = new ConcreteNode(this.m_heapctr++, tai, TypeIdentifier.OBJECT_HEADER_SIZE, TypeIdentifier.OBJECT_HEADER_SIZE);
			this.m_heap.AddNewObject(rr);
			this.m_visited.Add(o, rr);
			rr.FinalizeTargetSet(new List<ConcretePointer>());

			return rr;
		}
	}
}
